import quafu
from quafu import QuantumCircuit, simulate
import matplotlib.pyplot as plt

# ===== 调试开关 =====
SHOW_DEBUG_INFO = False  # 设为True显示过程信息

# ===== 创建量子电路 =====
qc = QuantumCircuit(2)
qc.h(0)
qc.cx(0, 1)

# ===== 状态向量模拟 =====
if SHOW_DEBUG_INFO:
    print("\n[调试] 开始状态模拟...")
state = simulate(qc).get_statevector()
print("\n=== 状态向量 ===")
print(state)

# ===== 测量结果获取 =====
qc_measure = QuantumCircuit(2)
qc_measure.h(0)
qc_measure.cx(0, 1)
qc_measure.measure([0])
qc_measure.measure([1])

if SHOW_DEBUG_INFO:
    print("\n[调试] 开始测量模拟...")
counts = simulate(qc_measure, shots=1000).counts
print("\n=== 测量结果 ===")
print(counts)

# ===== 可视化函数 =====
def plot_circuit():
    """带调试控制的绘图函数"""
    fig, ax = plt.subplots(figsize=(6, 3))
    
    # 固定坐标设置
    ax.set_xlim(-1, 3)
    ax.set_ylim(-0.5, 1.5)
    ax.set_axis_off()
    
    # 量子线绘制
    ax.hlines(0, -0.5, 2.5, color='black', lw=2)
    ax.hlines(1, -0.5, 2.5, color='black', lw=2)
    ax.text(-0.8, 0, 'q[0]', ha='right', va='center')
    ax.text(-0.8, 1, 'q[1]', ha='right', va='center')
    
    # H门绘制
    ax.text(0, 0, 'H', 
           ha='center', va='center',
           fontsize=20,
           bbox=dict(boxstyle="circle", fc="red", ec="black"))
    
    # CX门绘制
    ax.plot([2, 2], [0, 1], color='black', lw=2)
    ax.scatter(2, 0, s=200, facecolor='white', edgecolor='black')
    ax.text(2, 1, 'X', 
           ha='center', va='center',
           fontsize=20,
           bbox=dict(boxstyle="square", fc="white", ec="black"))
    
    if SHOW_DEBUG_INFO:
        print("[调试] 正在渲染图形...")
    plt.show()

# ===== 执行程序 =====
if __name__ == "__main__":
    print("\n=== 量子电路模拟 ===")
    plot_circuit()
    print("=== 模拟完成 ===")